package mr

import (
	"6.824/raft"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
)

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

type worker struct {
	id      int
	mapf    func(string, string) []KeyValue
	reducef func(string, []string) string
}

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	// uncomment to send the Example RPC to the coordinator.
	// CallExample()
	w := worker{}
	w.mapf = mapf
	w.reducef = reducef
	w.register()
	w.run()
}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	call("Coordinator.Example", &args, &reply)

	// reply.Y should be 100.
	raft.DPrintf("reply.Y %v\n", reply.Y)
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}

func (w *worker) doTask(t Task) {
	raft.DPrintf("in do Task %+v", t)
	switch t.Phase {
	case MapPhase:
		w.doMapTask(t)
	case ReducePhase:
		w.doReduceTask(t)
	default:
		panic(fmt.Sprintf("task phase err: %v", t.Phase))
	}
	raft.DPrintf("in do Task success\n")
}

func (w *worker) doMapTask(t Task) {
	raft.DPrintf("doMapTask exe\n")
	content, err := ioutil.ReadFile(t.FileName)
	if err != nil {
		w.reportTask(t, false)
	}
	kva := w.mapf(t.FileName, string(content))
	intermediate := make([][]KeyValue, t.NReduce)
	for _, kv := range kva {
		idx := ihash(kv.Key) % t.NReduce
		intermediate[idx] = append(intermediate[idx], kv)
	}
	for index, v := range intermediate {
		if len(v) == 0 {
			continue
		}
		filename := reduceName(t.Seq, index)
		f, err := ioutil.TempFile("", filename+"_tmp")
		if err != nil {
			w.reportTask(t, false)
		}
		enc := json.NewEncoder(f)
		for _, kv := range v {
			enc.Encode(&kv)
		}
		f.Close()
		err2 := os.Rename(f.Name(), filename)
		if nil != err2 {
			w.reportTask(t, false)
			log.Fatalf("map rename error %s  %+v", filename, err2)
		}
	}
	w.reportTask(t, true)
	raft.DPrintf("doMapTask success -- %s", t.FileName)
}

func (w *worker) doReduceTask(t Task) {
	raft.DPrintf("%d doReduceTask exe\n", t.Seq)
	maps := make(map[string][]string)

	for i := 0; i < t.NMaps; i++ {
		filename := reduceName(i, t.Seq)
		file, err := os.Open(filename)
		if err != nil {
			continue
		}
		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			if _, ok := maps[kv.Key]; !ok {
				maps[kv.Key] = make([]string, 0, 100)
			}
			maps[kv.Key] = append(maps[kv.Key], kv.Value)
		}
	}
	f, err := ioutil.TempFile("", "mr-out-tmp")
	if err != nil {
		w.reportTask(t, false)
	}
	for k, v := range maps {
		fmt.Fprintf(f, "%v %v\n", k, w.reducef(k, v))
	}

	err2 := os.Rename(f.Name(), mergeName(t.Seq))
	if nil != err2 {
		w.reportTask(t, false)
		raft.DPrintf("rename error %s  %+v", f.Name(), err2)
	}
	w.reportTask(t, true)
	raft.DPrintf("%d doReduceTask success -- %s", t.Seq, t.FileName)
}

func (w *worker) run() {
	for {
		t := w.reqTask()
		if !t.Alive {
			raft.DPrintf("worker get task not alive, exit task: %+v", t)
			break
		}
		raft.DPrintf("worker get task success task: %+v", t)
		w.doTask(t)
	}
}

func (w *worker) register() {
	raft.DPrintf("register exe..\n")

	args := &RegisterArgs{}
	rep := &RegisterReply{}
	if ok := call("Coordinator.RegisterWorker", &args, &rep); !ok {
		log.Fatalln("wroker register fail")
		os.Exit(1)
	}
	w.id = rep.WorkerId
	raft.DPrintf("register success rep : %+v", rep)
}

func (w *worker) reqTask() Task {
	raft.DPrintf("reqTask exe")
	args := TaskArgs{}
	args.WorkId = w.id
	reply := TaskReply{}
	if ok := call("Coordinator.GetOneTask", &args, &reply); !ok {
		log.Fatalln("wroker get task fail")
		os.Exit(1)
	}
	raft.DPrintf("worker get task:%+v", reply.Task)
	return *reply.Task
}

func (w *worker) reportTask(t Task, done bool) {
	args := ReportTaskArgs{}
	args.Done = done
	args.Seq = t.Seq
	args.Phase = t.Phase
	args.WorkerId = w.id
	args.FileName = t.FileName
	reply := RegisterReply{}
	if ok := call("Coordinator.ReportTask", &args, &reply); !ok {
		raft.DPrintf("reportTask task error - %+v", t.Seq)
	} else {
		raft.DPrintf("reportTask task success - %+v", t.Seq)
	}
}

func reduceName(mapIdx, reduceIdx int) string {
	return fmt.Sprintf("mr-%d-%d", mapIdx, reduceIdx)
}

func mergeName(reduceIdx int) string {
	return fmt.Sprintf("mr-out-%d", reduceIdx)
}
